#!/usr/bin/env python3
# -*-coding:utf8-*-

import json
import logging
import os

os.environ["MUJOCO_GL"] = "egl"
import time
from pathlib import Path
from typing import Dict, List, Sequence

import numpy as np
from piper_sdk import C_PiperInterface_V2
from scipy.spatial.transform import Rotation

from sim.piper_simulator import MujocoPiperSim

"""
V1: move to position using MujocoIK, abandon the interal IK:
PosEndCtrl is never used

"""


log = logging.getLogger(__name__)

JOINT_LIMITS = np.array([[-2.6179, 0, -2.967, -1.745, -1.22, -2.09439], [2.6179, 3.14, 0, 1.745, 1.22, 2.09439]])


def radian_to_machine_angle(joints: np.ndarray):
    """
    radian -> 0.001 degree
    """
    return (np.rad2deg(joints) * 1000).astype(int)


def machine_angle_to_radian(joints: np.ndarray):
    """
    0.001 degree -> radian
    """
    return np.deg2rad(joints) / 1000


def clip_joint_radian(joints: np.ndarray):
    """
    clip the joint radian to the working range of the robot
    """
    return np.clip(joints, JOINT_LIMITS[0], JOINT_LIMITS[1])


class PiperArm:
    """
    Unify the measurement unit:
    Length: mm
    Angle: radian
    Time: s
    Speed: mm/s

    __speed: mm/s (int) gives the running speed of the arm
    __Hz: Hz (int) gives the frequency of the signal sent to the arm
    __timeout: s (int) gives the timeout of the arm execution
    __gripper_effort: 0~5000, 0~5N/m
    __gripper_angle: 0~10000, 0~100°

    Caution:
        ALL actions are executed in JOINT mode through "move_to_joint" method
        The joint control = piper movement + ik update + wait_motion_done

    """

    __Hz: int = 100
    __timeout: int = 5
    __gripper_effort: int = 1000
    __wait_gripper: float = 0.3
    speed: int = 100

    def __init__(
        self,
        can: str = "can0",
        wait_execution: bool = True,
    ):  # in mm
        self.ik = MujocoPiperSim()

        self.__timestep = 1 / self.__Hz
        self.can = can
        self.wait_execution = wait_execution

        # Initialize the robot
        self.arm = C_PiperInterface_V2(self.can)
        self.arm.ConnectPort()
        if self.reset():
            print(f"Sucessfully connect to the robot CAN={can}")
            print(f"Sucessfully make arm {can} enabled")

        # table position transformation
        self.table_M = None
        if os.path.exists(f"data/table_calibration_{self.can}.json"):
            with open(f"data/table_calibration_{self.can}.json", "r") as f:
                self.table_M = np.array(json.load(f))
            print(f"Successfully loaded table calibration matrix from data/table_calibration_{self.can}.json")

    def reset(self):
        s_time = time.time()
        self.go_zero()
        while not all(self.arm.GetArmEnableStatus()):
            self.arm.EnableArm(7)
            self.go_zero()
            time.sleep(self.__timestep)
            if time.time() - s_time > self.__timeout:
                log.error("Arm reset timeout!")
                return False
        return True

    def go_zero(self):
        """
        use joint mode to move to the zero position
        """
        self.arm.ModeCtrl(0x01, 0x01, self.speed, 0x00)
        self.arm.JointCtrl(0, 0, 0, 0, 0, 0)
        self.arm.GripperCtrl(0, self.__gripper_effort, 0x01, 0)
        if self.wait_execution:
            self.wait_motion_done()

    def go_home(self):
        self.move_to_joint(self.ik.j_home[:6])

    def get_rpy_prior(self, xyz: np.ndarray, gripper_slope: float = 10.0):
        # large gripper_slope means the gripper is more vertical
        ry = min(gripper_slope * np.arctan2(1.0, np.sqrt(xyz[0] ** 2 + xyz[1] ** 2)) + np.pi / 2, np.pi)
        rz = np.arctan2(xyz[1], xyz[0])
        return np.array([0, ry, rz])

    def move_to(self, xyz: np.ndarray, rpy: np.ndarray = None, down_to_table: bool = False):
        """
        USE THIS TO REACH CONTACT POINT !!!!

        move to the contact position (gripper)
        xyz (mm) : the position of the end-effector

        The process:
        1. clip the xyz to the same zone
        2. If ik model loaded: predict the joint state -> joint move
        3. If no ik model: estimate the arm position -> arm move
        """
        assert len(xyz) == 3, "xyz should be a 3D vector"
        assert max(np.abs(xyz)) < 1, "the unite of xyz positions is meter!"

        if down_to_table:
            xyz = self.table_pos_from_xyz(xyz)

        if rpy is None:
            # the gripper is vertical to the table
            rpy = self.get_rpy_prior(xyz)

        # visualization
        self.ik.move_target_xeuler(xyz, rpy)

        prev_j = np.concat([self.get_joint_status(), self.get_gripper_status()])
        joints = self.ik.solve_ik_xyz_rpy(xyz, rpy, j0=None, x_prev=prev_j.copy())
        self.move_to_joint(joints)

    def planning_to(self, planning_model, pos: np.ndarray, rpy: np.ndarray, render: bool = False):

        self.ik.move_target_xeuler(pos, rpy)
        self.step_ik()
        cond = np.concat([self.ik.get_effector_xeuler(), pos, rpy], axis=-1)
        # (batch_size, num_act_samples, horizon_steps, xeuler=6)
        path = planning_model(cond)[0, 0]
        if render:
            frames = [self.ik.render_frame()]
        else:
            frames = []
        for pt in path:
            self.move_to(pt[:3], pt[3:])
            if render:
                frames.append(self.ik.render_frame())
        return frames

    def step_ik(self):
        """
        update the ik model with the current joint status of the real arm
        """
        j = np.concat([self.get_joint_status(), self.get_gripper_status()])
        self.ik.step(j)

    def render_ik(self):
        self.step_ik()
        self.ik.render()

    def open_gripper(self, width: int = 100):
        """
        width: 0~100
        need some time to finish the operation: self.__wait_gripper
        """
        self.wait_motion_done()
        self.arm.GripperCtrl(width * 1000, self.__gripper_effort, 0x03, 0)
        time.sleep(self.__wait_gripper)

    def close_gripper(self):
        """
        need some time to finish the operation: self.__wait_gripper
        """
        self.wait_motion_done()
        self.arm.GripperCtrl(0, self.__gripper_effort, 0x03, 0)
        time.sleep(self.__wait_gripper)

    def move_to_joint(self, joints: np.ndarray):
        """
        joints: radian -> clip -> 0.001 degree (int) -> machine execution
        """

        # joints = clip_joint_radian(joints[:6])
        machine_joints = radian_to_machine_angle(clip_joint_radian(joints[:6]))

        # "Joint" mode
        self.arm.ModeCtrl(0x01, 0x01, self.speed, 0x00)
        self.arm.JointCtrl(*machine_joints)

        if self.wait_execution:
            return self.wait_motion_done()

        self.ik.step(joints)

    def get_gripper_status(self):
        return (
            np.array([1, -1])
            * machine_angle_to_radian(self.arm.GetArmGripperMsgs().gripper_state.grippers_angle).item()
            / 1.78
            * 0.04
        )

    def get_joint_status(
        self,
    ) -> Dict:
        arm_status = self.arm.GetArmJointMsgs().joint_state
        jo = [
            arm_status.joint_1,
            arm_status.joint_2,
            arm_status.joint_3,
            arm_status.joint_4,
            arm_status.joint_5,
            arm_status.joint_6,
        ]
        return machine_angle_to_radian(np.asarray(jo))

    def get_effector_xeuler(self):
        self.step_ik()
        return self.ik.get_effector_xeuler()

    def is_in_motion(self):
        return bool(self.arm.GetArmStatus().arm_status.motion_status)

    def wait_motion_done(self):
        """
        Either the effector is close to the target position or the arm is not in motion
        """
        start_t = time.time()
        time.sleep(self.__timestep)
        while self.ik.et_error(self.get_joint_status() + 1e-6) > 0.01:  # 1cm tolerance
            time.sleep(self.__timestep)
            if not self.is_in_motion():
                return True
            if time.time() - start_t > self.__timeout:
                log.error(f"Excution exceeds timeout threshold={self.__timeout}s")
                return False
        return True

    def is_gripping_sth(self):
        return self.arm.GetArmGripperMsgs().gripper_state.grippers_effort < -200

    def no_wait_execution(self):
        """
        temporary disable the wait_execution
        with arm.no_wait_execution():
            print(arm.wait_execution)  # False
        """

        class NoWaitContext:
            def __init__(self, arm):
                self.arm = arm
                self.original_wait_execution = None

            def __enter__(self):
                self.original_wait_execution = self.arm.wait_execution
                self.arm.wait_execution = False
                return self.arm

            def __exit__(self, exc_type, exc_val, exc_tb):
                self.arm.wait_execution = self.original_wait_execution

        return NoWaitContext(self)

    def table_pos_calibration(self):
        """
        Use linear regression to fit the table surface:
        z = ax + by + cr + d
        r = sqrt(x^2 + y^2)
        """
        print("Please place the gripper on the table and activate the teach mode")
        key = input("Press ENTER to start recording...")
        xy, z = [], []
        try:
            while True:
                while self.arm.GetArmStatus().arm_status.teach_status == 1:
                    xyz = self.get_effector_xeuler()[:3]
                    xy.append(xyz[:2].tolist())
                    z.append(xyz[2].item())
                    time.sleep(0.1)
                    print(f"Recorded {len(xy)} table surface points", end="\r")
                time.sleep(0.1)

        except KeyboardInterrupt:
            if len(xy) < 3:
                print("Not enough points to compute table calibration matrix")
                return

            xy_array = np.array(xy)
            z_array = np.array(z)
            r = np.sqrt(xy_array[:, 0] ** 2 + xy_array[:, 1] ** 2)
            xyr = np.hstack([xy_array, r[:, None]])
            A_hom = np.hstack([xyr, np.ones((len(xyr), 1))])
            M, residuals, rank, singular = np.linalg.lstsq(A_hom, z_array, rcond=None)
            print(f"Computed table calibration matrix: {M} (residuals={residuals})")
            with open(f"data/table_calibration_{self.can}.json", "w") as f:
                json.dump(M.tolist(), f)  # 1x4 matrix describes z = ax + by + cr + d
            M[3] -= residuals  # use lower bound
            self.table_M = M

        # visualization of table xyz surface
        from matplotlib import pyplot as plt
        from mpl_toolkits.mplot3d import Axes3D

        fig = plt.figure(figsize=(12, 8))
        ax = fig.add_subplot(111, projection="3d")

        ax.scatter(xy_array[:, 0], xy_array[:, 1], z_array * 100, c="red", marker="o", s=50, label="table data")

        # create grid for surface
        x_min, x_max = xy_array[:, 0].min(), xy_array[:, 0].max()
        y_min, y_max = xy_array[:, 1].min(), xy_array[:, 1].max()

        # expand grid range for better surface display
        x_range = x_max - x_min
        y_range = y_max - y_min
        x_min -= x_range * 0.1
        x_max += x_range * 0.1
        y_min -= y_range * 0.1
        y_max += y_range * 0.1

        # create grid points
        x_grid, y_grid = np.meshgrid(np.linspace(x_min, x_max, 50), np.linspace(y_min, y_max, 50))
        r_grid = np.sqrt(x_grid**2 + y_grid**2)

        # compute z value on the surface: z = ax + by + c
        z_grid = M[0] * x_grid + M[1] * y_grid + M[2] * r_grid + M[3]

        # plot surface
        surf = ax.plot_surface(x_grid, y_grid, z_grid * 100, alpha=0.6, color="blue", label="surface")

        # set graph properties
        ax.set_xlabel("X (m)")
        ax.set_ylabel("Y (m)")
        ax.set_zlabel("Z (cm)")
        ax.set_title(
            f"Table surface calibration - regression equation: z = {M[0]:.4f}x + {M[1]:.4f}y + {M[2]:.4f}r + {M[3]:.4f}"
        )

        # add legend
        ax.view_init(elev=10, azim=120)
        ax.legend()

        # show graph
        plt.tight_layout()
        plt.show()

        plt.figure()
        plt.scatter(r, z_array * 100, label="Z (cm)")
        plt.xlabel("r (m)")
        plt.ylabel("Z (cm)")
        plt.title("Z vs r")
        plt.legend()
        plt.show()

        return M

    def table_pos_from_xyz(self, xyz: np.ndarray):
        if self.table_M is not None:
            r = np.sqrt(xyz[0] ** 2 + xyz[1] ** 2)
            z = np.append(np.append(xyz[:2], r), 1) @ self.table_M
            return np.append(xyz[:2], z)
        else:
            print("Table calibration not found, please run table_pos_calibration first")
            return xyz

    def follow_traj(self, poss: np.ndarray, rpys: np.ndarray = None):
        """
        Only support joint mode
        """
        for i in range(len(poss)):
            if rpys is None:
                self.move_to(poss[i])
            else:
                self.move_to(poss[i], rpys[i])

    def record_traj(self, filename: Path = "data/json_data/trajectories.json"):
        xeulers, joints, lengths = [], [], []
        try:
            while True:
                xe, j = [], []
                while self.arm.GetArmStatus().arm_status.teach_status == 1:
                    xe.append(self.get_effector_xeuler().tolist())
                    j.append(self.get_joint_status().tolist())
                    time.sleep(0.1)
                    print(f"Recorded {len(xe)} points", end="\r")

                if xeulers and joints:
                    xeulers += xe
                    joints += j
                    lengths.append(len(xeulers))
                    print(f"Recorded {len(xeulers)} points: Done!", end="\n")
                else:
                    print("No points recorded, please move the arm to record...", end="\r")
                    time.sleep(0.1)
        except KeyboardInterrupt:
            with open(filename, "w") as f:
                json.dump(
                    {
                        "xeuler": xeulers,
                        "joint": joints,
                        "length": lengths,
                    },
                    f,
                )

    def record_points(
        self, save_dir: Path = "data/json_data", method: str = "auto", tag: str = "", time_interval: float = 0.02
    ):
        """
        record the joint and ee position status
        """
        n_ = 0
        joints, xeulers = [], []

        if not os.path.exists(save_dir):
            os.mkdir(save_dir)

        try:
            while True:
                if method != "auto":
                    key = input("Place the gripper and press ENTER...")
                    if key == "q":
                        assert 1 == 0
                joints.append(self.get_joint_status().tolist())
                xeulers.append(self.get_effector_xeuler().tolist())
                n_ += 1

                time.sleep(time_interval)
                print(f"Saved {n_} points", end="\r")

        except KeyboardInterrupt:

            f_name = os.path.join(save_dir, f"points_{tag}.json")
            with open(f_name, "w") as f:
                json.dump(
                    {
                        "joint": joints,
                        "xeuler": xeulers,
                    },
                    f,
                )
                print(f"{n_} data points are saved to {f_name}")

    def mujoco_viewer(self, run_in_background=True):
        """
        Used for notebook running
        """
        os.environ["MUJOCO_GL"] = "egl"
        import threading

        import mujoco
        import mujoco.viewer

        model = mujoco.MjModel.from_xml_path("piper_description/mujoco_model/piper_description.xml")
        data = mujoco.MjData(model)

        def _viewer_loop():
            print("Viewer launced!")
            with mujoco.viewer.launch_passive(model, data) as viewer:
                print("Viewer launched!")
                while viewer.is_running():
                    joint = np.concat([self.get_joint_status(), self.get_gripper_status()])
                    # joint[1] += np.pi /90
                    data.qpos = joint
                    # TODO: use data.ctrl to control the robot
                    mujoco.mj_step(model, data)
                    viewer.sync()
                    time.sleep(0.01)
                print("Viewer closed!")

        if run_in_background:
            # daemon=True: thread will be terminated when the main thread terminates
            self.viewer_thread = threading.Thread(target=_viewer_loop, daemon=True)
            self.viewer_thread.start()
            print("Viewer started in background thread. You can now run other cells!")
            return self.viewer_thread
        else:
            _viewer_loop()


if __name__ == "__main__":

    def test_planning():
        import hydra
        from omegaconf import DictConfig

        with hydra.initialize(version_base=None, config_path="./cfg"):
            ej_cfg: DictConfig = hydra.compose(
                config_name="train_diffusion_planning",
                overrides=[
                    "tag=10K",
                    "model.temperature=0.1",
                    "model.horizon_steps=10",
                    "model.num_act_samples=100",
                ],
            )
            model = hydra.utils.instantiate(ej_cfg.model)
            model.load_ckpt(ej_cfg.ckp_name + ".ckp")

        self = PiperArm("can1", wait_execution=True)
        pos1 = np.array([0.408, -0.068, 0.045])
        pos2 = np.array([0.008, -0.468, 0.045])

        rpy = np.array([0, 3.1, 0])

        mid_pos = np.linspace(pos1, pos2, 100)
        # for i in range(10):

        thetas = np.linspace(0, 2 * np.pi, 500)
        xyz = np.stack([np.cos(thetas), np.sin(thetas), np.zeros_like(thetas)], axis=-1)
        circle = xyz * 0.2 + np.array([0.4, 0, 0])
        self.follow_traj(circle)

        self.move_to(pos1)
        self.follow_traj(mid_pos)
        self.wait_motion_done()
        self.render_ik()
        # self.planning_to(model, pos=pos, rpy=rpy)
